package middleware

import (
	"bytes"
	"context"
	"github.com/gin-gonic/gin"
	"io"
	"time"
)

type LogMiddlewareBuilder struct {
	logFunc func(ctx context.Context, l AccessLog)
	//线上环境最好不要打印req和resp 防止攻击者恶意攻击
	allowReqBody  bool
	allowRespBody bool
}

func NewLogMiddlewareBuilder(logFc func(ctx context.Context, l AccessLog)) *LogMiddlewareBuilder {
	return &LogMiddlewareBuilder{
		logFunc: logFc,
	}
}
func (l *LogMiddlewareBuilder) AllowReqBody() *LogMiddlewareBuilder {
	l.allowReqBody = true
	return l
}

func (l *LogMiddlewareBuilder) AllowRespBody() *LogMiddlewareBuilder {
	l.allowRespBody = true
	return l
}
func (l LogMiddlewareBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		start := time.Now()
		path := ctx.Request.URL.Path
		if len(path) > 1024 { //防止攻击者
			path = path[:1024]
		}
		method := ctx.Request.Method
		al := AccessLog{
			Path:   path,
			Method: method,
		}
		if l.allowReqBody {
			body, _ := ctx.GetRawData()
			if len(body) > 2048 {
				al.ReqBody = string(body[:2048])
			} else {
				al.ReqBody = string(body)
			}
			//request的body是流对象 只能读取一次 因为读完要放回去 不然后续步骤读不到
			ctx.Request.Body = io.NopCloser(bytes.NewBuffer(body))
		}
		if l.allowRespBody {
			ctx.Writer = responseWriter{
				ResponseWriter: ctx.Writer,
				al:             &al,
			}
		}
		defer func() {
			duration := time.Since(start)
			al.Duration = duration.String()
			l.logFunc(ctx, al)
		}()
		ctx.Next() //继续执行业务代码

		//在这里拿到了响应

	}
}

type AccessLog struct {
	Path       string `json:"path"`
	Method     string `json:"method"`
	ReqBody    string `json:"req_body"`
	RespBody   string `json:"resp_body"`
	StatusCode int    `json:"status_code"`
	Duration   string `json:"duration"`
}

type responseWriter struct {
	al *AccessLog
	gin.ResponseWriter
}

func (r responseWriter) WriteHeader(statusCode int) {
	r.al.StatusCode = statusCode
	r.ResponseWriter.WriteHeader(statusCode)
}

func (r responseWriter) Write(data []byte) (int, error) {
	r.al.RespBody = string(data)
	return r.ResponseWriter.Write(data)
}

func (r responseWriter) WriteString(data string) (int, error) {
	r.al.RespBody = data
	return r.ResponseWriter.WriteString(data)
}
